"""
照片名字相似度分类模块
用于将名字格式相似的照片自动分类到同一文件夹
"""

import os
import re
from difflib import SequenceMatcher
from collections import defaultdict


class NameClassifier:
    """照片名字分类器"""

    def __init__(self, similarity_threshold=0.6):
        """
        初始化分类器

        Args:
            similarity_threshold: 相似度阈值，范围0-1，默认0.6
                                 值越大要求越严格，越小则越宽松
        """
        self.similarity_threshold = similarity_threshold

    def extract_pattern(self, filename):
        """
        提取文件名的核心结构模式
        保留商品名称和编号，替换颜色和尾号等变化部分

        Args:
            filename: 文件名（不含扩展名）

        Returns:
            提取的模式字符串
        """
        # 移除文件扩展名（如果有）
        name = os.path.splitext(filename)[0]

        # 第一步：移除结尾的 -数字 (这通常是照片序号)
        pattern = re.sub(r'-\d+$', '', name)

        # 第二步：识别并替换颜色相关的词
        # 扩展的颜色词库，包括常见颜色和修饰词
        color_keywords = [
            # 基础颜色
            '红', '橙', '黄', '绿', '青', '蓝', '紫', '黑', '白', '灰', '棕', '咖', '褐',
            # 特殊颜色
            '米', '粉', '杏', '驼', '卡其', '军绿', '藏青', '宝蓝', '天蓝', '湖蓝', '孔雀蓝',
            '玫红', '酒红', '枣红', '大红', '砖红', '姜黄', '鹅黄', '柠檬黄',
            '墨绿', '草绿', '橄榄绿', '浅绿', '深绿',
            '浅蓝', '深蓝', '藏蓝', '宝蓝',
            '浅灰', '深灰', '炭灰', '烟灰',
            '米白', '象牙白', '乳白', '奶白', '本白',
            '深咖', '浅咖', '咖啡',
            # 颜色后缀
            '色',
        ]

        # 创建一个正则表达式，匹配所有颜色词的组合
        # 例如：深灰色、红咖色、白蓝色等
        color_pattern = '|'.join(re.escape(word) for word in sorted(color_keywords, key=len, reverse=True))

        # 匹配连续的颜色词组合（如"深灰色"、"白蓝色"）
        pattern = re.sub(f'({color_pattern})+', '[COLOR]', pattern)

        return pattern

    def calculate_similarity(self, name1, name2):
        """
        计算两个文件名的相似度
        基于去除颜色和尾号后的核心结构进行比较

        Args:
            name1: 第一个文件名
            name2: 第二个文件名

        Returns:
            相似度分数，范围0-1
        """
        # 提取两个文件名的模式
        pattern1 = self.extract_pattern(name1)
        pattern2 = self.extract_pattern(name2)

        # 如果模式完全相同，直接返回1.0
        if pattern1 == pattern2:
            return 1.0

        # 计算模式的相似度
        similarity = SequenceMatcher(None, pattern1, pattern2).ratio()

        return similarity

    def group_by_similarity(self, filenames):
        """
        根据相似度将文件名分组

        Args:
            filenames: 文件名列表

        Returns:
            分组字典，{组ID: [文件名列表]}
        """
        if not filenames:
            return {}

        groups = []
        ungrouped = filenames.copy()

        while ungrouped:
            # 取第一个未分组的文件作为新组的代表
            current = ungrouped[0]
            current_group = [current]
            ungrouped.remove(current)

            # 查找与当前文件相似的其他文件
            to_remove = []
            for filename in ungrouped:
                similarity = self.calculate_similarity(current, filename)
                if similarity >= self.similarity_threshold:
                    current_group.append(filename)
                    to_remove.append(filename)

            # 从未分组列表中移除已分组的文件
            for filename in to_remove:
                ungrouped.remove(filename)

            groups.append(current_group)

        # 转换为字典格式
        result = {}
        for i, group in enumerate(groups):
            result[i] = group

        return result

    def generate_folder_name(self, filenames, max_length=80):
        """
        为一组文件生成文件夹名称
        使用去除颜色和尾号后的核心结构

        Args:
            filenames: 文件名列表
            max_length: 文件夹名最大长度

        Returns:
            生成的文件夹名
        """
        if not filenames:
            return "未命名分组"

        # 使用第一个文件名的模式（去掉颜色）作为文件夹名
        first_file = os.path.splitext(filenames[0])[0]
        pattern = self.extract_pattern(first_file)

        # 移除[COLOR]占位符，让名称更简洁
        folder_name = pattern.replace('[COLOR]', '').replace('++', '+')

        # 清理多余的符号
        folder_name = re.sub(r'\+\s*\+', '+', folder_name)  # 去掉重复的+
        folder_name = folder_name.strip('+').strip()  # 去掉首尾的+和空格

        # 如果名称太短或为空，使用原始名称
        if len(folder_name) < 5:
            folder_name = first_file

        # 截断过长的名称
        if len(folder_name) > max_length:
            folder_name = folder_name[:max_length] + "..."

        # 清理不合法的文件夹字符
        folder_name = self.sanitize_folder_name(folder_name)

        # 添加组内文件数量
        folder_name = f"{folder_name} ({len(filenames)}张)"

        return folder_name

    def sanitize_folder_name(self, name):
        """
        清理文件夹名称中的非法字符

        Args:
            name: 原始名称

        Returns:
            清理后的名称
        """
        # Windows文件夹名不允许的字符
        illegal_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*']

        for char in illegal_chars:
            name = name.replace(char, '_')

        # 移除首尾空格
        name = name.strip()

        # 如果名称为空，返回默认名称
        if not name:
            name = "新建分组"

        return name

    def classify_photos(self, photo_folder, preview_only=False):
        """
        对照片文件夹中的照片进行分类

        Args:
            photo_folder: 照片所在文件夹路径
            preview_only: 仅预览不实际创建文件夹

        Returns:
            分类结果字典 {文件夹名: [文件列表]}
        """
        # 获取所有图片文件
        valid_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp']

        photos = []
        for file in os.listdir(photo_folder):
            file_path = os.path.join(photo_folder, file)
            if os.path.isfile(file_path):
                ext = os.path.splitext(file)[1].lower()
                if ext in valid_extensions:
                    photos.append(file)

        if not photos:
            return {}

        # 根据相似度分组
        groups = self.group_by_similarity(photos)

        # 为每组生成文件夹名
        classification_result = {}
        for group_id, filenames in groups.items():
            # 只有当组内文件数大于1时才创建分组
            if len(filenames) > 1:
                folder_name = self.generate_folder_name(filenames)
                classification_result[folder_name] = filenames

        return classification_result

    def create_classified_folders(self, photo_folder, classification_result, move_files=False):
        """
        根据分类结果创建文件夹并移动/复制文件

        Args:
            photo_folder: 照片所在文件夹
            classification_result: 分类结果字典
            move_files: True为移动文件，False为复制文件

        Returns:
            成功创建的文件夹数量
        """
        import shutil

        created_count = 0

        for folder_name, filenames in classification_result.items():
            # 创建目标文件夹
            target_folder = os.path.join(photo_folder, folder_name)

            try:
                # 如果文件夹不存在则创建
                if not os.path.exists(target_folder):
                    os.makedirs(target_folder)

                # 移动或复制文件
                for filename in filenames:
                    source_path = os.path.join(photo_folder, filename)
                    dest_path = os.path.join(target_folder, filename)

                    # 跳过已经在目标文件夹的文件
                    if source_path == dest_path:
                        continue

                    # 如果目标文件已存在，添加序号
                    if os.path.exists(dest_path):
                        base, ext = os.path.splitext(filename)
                        i = 1
                        while os.path.exists(dest_path):
                            new_filename = f"{base}_{i}{ext}"
                            dest_path = os.path.join(target_folder, new_filename)
                            i += 1

                    # 移动或复制文件
                    if move_files:
                        shutil.move(source_path, dest_path)
                    else:
                        shutil.copy2(source_path, dest_path)

                created_count += 1

            except Exception as e:
                print(f"创建文件夹 {folder_name} 时出错: {str(e)}")
                continue

        return created_count


def test_classifier():
    """测试分类器功能"""
    classifier = NameClassifier(similarity_threshold=0.95)

    # 使用真实的用户示例进行测试
    test_names = [
        "不对称多穿马海毛中长针织上衣363191142蓝色+不对称搭片羊毛A字裤364250729深灰色-9.jpg",
        "不对称多穿马海毛中长针织上衣363191142深灰色+不对称搭片羊毛A字裤364250729黑色-4.jpg",
        "反折搭片多穿针织上衣364291720白蓝色+后开衩毛边羊毛半裙364154172红褐色-5.jpg",
        "反折搭片多穿针织上衣364291720黑色+后开衩毛边羊毛半裙364154172黑色-4.jpg",
        "弧形不对称宽松多穿针织上衣363110161米白色+搭片双层A字半裙363154162米白色-5.jpg",
        "弧形不对称宽松多穿针织上衣363110161米白色+搭片双层A字半裙363154162米白色-8.jpg",
    ]

    print("=" * 80)
    print("测试文件名分组:")
    print("=" * 80)

    # 打印每个文件的模式
    print("\n【步骤1】提取每个文件的核心模式:")
    for i, name in enumerate(test_names, 1):
        pattern = classifier.extract_pattern(name)
        print(f"{i}. 原文件名: {name}")
        print(f"   核心模式: {pattern}\n")

    # 进行分组
    groups = classifier.group_by_similarity(test_names)

    print("=" * 80)
    print(f"\n【步骤2】分组结果（共 {len(groups)} 组）:")
    print("=" * 80)

    for group_id, filenames in groups.items():
        print(f"\n组 {group_id + 1} - 共 {len(filenames)} 个文件:")
        folder_name = classifier.generate_folder_name(filenames)
        print(f"文件夹名: {folder_name}\n")
        for filename in filenames:
            print(f"  - {filename}")

    print("\n" + "=" * 80)
    print("测试完成！")
    print("=" * 80)


if __name__ == "__main__":
    test_classifier()
